import json
import time

import cv2

import t_yolov5v6_2onnx

# 类别
names=['front_wear', 'front_no_wear', 'front_under_nose_wear', 'front_under_mouth_wear', 'front_unknown', 'side_wear', 'side_no_wear', 'side_under_nose_wear', 'side_under_mouth_wear', 'side_unknown', 'side_back_head_wear', 'side_back_head_no_wear', 'back_head', 'mask_front_wear', 'mask_front_under_nose_wear', 'mask_front_under_mouth_wear', 'mask_side_wear', 'mask_side_under_nose_wear', 'mask_side_under_mouth_wear', 'strap']


def init():

    # Initialize
    modelpath_car =  r'/project/train/models/exp/weights/best.onnx'
    # 加载 模型
    yolonet_model = t_yolov5v6_2onnx.yolov5(modelpath_car, confThreshold=0.3, iou_thres=0.5) #刚开始就限定了，置信率
    model=yolonet_model

    return model

def process_image(model, input_image=None, args=None, **kwargs):

    srcimg, list_bboxs = model.detect(input_image)  # 进行 图片 推理

    cv2.imshow('srcimg_car', srcimg)
    cv2.waitKey(show_time)

    fake_result = {}

    fake_result["algorithm_data"] = { #也就是说，这个这里 只是带上的，不要管
       "is_alert": False,
       "target_count": 0,
       "target_info": []
   }

    fake_result["model_data"] = {"objects": []}

    # Process detections
    cnt = 0

    for box in list_bboxs:
        #box = list_bboxs[i]
        left = box[0]
        top = box[1]

        width = box[2]
        height = box[3]
        confidence=float(box[5])
        name = box[4]

        cnt += 1
        fake_result["model_data"]['objects'].append({  # 这里 才是，真正要添加的
                "x": left,
                "y": top,
                "height": height,
                "width": width,
                "confidence": confidence,
                "name": name
        })

        fake_result["algorithm_data"]["target_info"].append({
            "x": left,
            "y": top,
            "height": height,
            "width": width,
            "confidence": confidence,
            "name": name
        })

    if cnt:
        fake_result ["algorithm_data"]["is_alert"] = True
        fake_result ["algorithm_data"]["target_count"] = cnt
    return json.dumps(fake_result, indent = 4)



show_time=10000


if __name__ == '__main__':

    from glob import glob

    # modelpath_car = r'/project/train/models/exp/weights/best.onnx'
    # image_names = glob(r'/home/data/1568/*.jpg')


    image_names = glob(r'/home/deepin/Documents/ji_pingtai/windows/yolov5/datasets/images/valid/3.jpg')

    predictor = init()

    s = 0
    count=0
    for image_name in image_names:
        #print('image_path:', os.path.join(image_dir, image_name))
        img = cv2.imread(image_name)

        start = time.time()

        res = process_image(predictor, img)

        print(res)

        end = time.time()

        inf_end= end - start
        s += end - start

        fps = 1 / inf_end

        fps_label = "FPS: %.2f " % fps
        fps_label = fps_label + " %.2f" % inf_end + "s " + f'({((float(inf_end) * 1000.0)):.1f}ms) Inference'
        print(fps_label)

    print(1 / (s / 100))

